import os
import torch
import torch.distributed as dist
from torch.nn.parallel import DistributedDataParallel as DDP
from torch.utils.data.distributed import DistributedSampler
import time
import datetime
import json


class DDPUtil:
    """分布式训练工具"""

    
    @staticmethod
    def get_backend():
        """根据操作系统和可用性选择最佳后端"""
        # Windows上强制使用gloo，因为nccl在Windows上不支持
        if os.name == 'nt':
            print("Windows系统检测到，强制使用gloo后端")
            return 'gloo'  # Windows上直接返回gloo，不需要额外检查，因为如果不支持会在初始化时报错
        else:
            # Linux/macOS优先使用nccl
            if dist.is_nccl_available():
                print("Linux/macOS系统检测到，使用nccl后端")
                return 'nccl'
            elif dist.is_gloo_available():
                print("nccl不可用，回退到gloo后端")
                return 'gloo'
            else:
                return None
    
    @staticmethod
    def setup(rank, world_size):
        """初始化分布式环境
        
        Args:
            rank: 进程的本地排名
            world_size: 总进程数
            
        Returns:
            str: 使用的后端名称
            
        Raises:
            Exception: 如果初始化失败
        """

        # 清除特定的调试环境变量，但保留MASTER_ADDR和MASTER_PORT（如果已存在）
        for key in ['NCCL_DEBUG', 'TORCH_DISTRIBUTED_DEBUG']:
            if key in os.environ:
                del os.environ[key]
        
        # 优先使用环境变量中的值，若不存在则设置默认值
        # 这样可以支持torchrun自动分配的端口号
        if 'MASTER_ADDR' not in os.environ:
            os.environ['MASTER_ADDR'] = 'localhost'
        
        # 为每个进程动态分配不同的端口范围，避免冲突
        if 'MASTER_PORT' not in os.environ:
            # 基础端口号，进程0使用12355，进程1使用12356，以此类推
            base_port = 12355 + rank
            os.environ['MASTER_PORT'] = str(base_port)
            print(f"进程 {rank} 使用端口: {base_port}")
        
        # 获取可用的后端
        backend = DDPUtil.get_backend()
        if backend is None:
            raise RuntimeError("没有可用的分布式后端，请检查PyTorch安装")
        
        # 针对Linux环境的特殊优化配置
        if os.name != 'nt':  # Linux/Unix环境
            if backend == 'nccl':
                # 优化NCCL通信性能
                os.environ['NCCL_IB_DISABLE'] = '0'  # 启用InfiniBand(如果可用)
                os.environ['NCCL_SOCKET_IFNAME'] = 'eth0,ib0'  # 优先使用的网络接口
                os.environ['NCCL_DEBUG'] = 'INFO'
                os.environ['NCCL_ASYNC_ERROR_HANDLING'] = '1'
                os.environ['NCCL_IB_TIMEOUT'] = '23'
                
                # 增加NCCL初始化超时时间
                os.environ['NCCL_TIMEOUT'] = '300'
        
        else:
            os.environ['TORCH_DISTRIBUTED_DEBUG'] = 'INFO'
        
        start_time = time.time()
        
        try:

            # 对于nccl后端，添加额外的NCCL调试信息
            if backend == 'nccl':
                os.environ['NCCL_DEBUG'] = 'DETAIL'  # 使用更详细的NCCL调试级别
                os.environ['NCCL_DEBUG_SUBSYS'] = 'INIT,ENV,GRAPH'  # 关注初始化、环境和图构建

            # 为了避免无限卡住，使用单独的线程进行初始化并设置超时
            import threading
            import queue
            
            result_queue = queue.Queue()
            error_queue = queue.Queue()
            
            def init_thread():
                try:
                    # 为Windows环境添加特殊处理
                    if os.name == 'nt':
                        # Windows上使用文件系统作为store后端，更稳定
                        import tempfile
                        temp_dir = tempfile.gettempdir()
                        os.environ['TORCH_DDP_SHARED_FILESYSTEM'] = temp_dir
                        
                        # torch 1.10.2 兼容的初始化方式，移除timeout参数以兼容旧版PyTorch
                        dist.init_process_group(
                            backend=backend,
                            rank=rank,
                            world_size=world_size
                        )
                    else:
                        # Linux/macOS上的初始化
                        dist.init_process_group(
                            backend=backend,
                            rank=rank,
                            world_size=world_size,
                        )
                    result_queue.put(True)
                except Exception as e:
                    error_queue.put(e)
                    result_queue.put(False)
            
            # 启动初始化线程
            thread = threading.Thread(target=init_thread)
            thread.daemon = True  # 设置为守护线程，主进程退出时自动终止
            thread.start()
            
            # 设置合理的超时时间，根据进程数调整
            timeout = 20.0 if world_size > 2 else 15.0
            thread.join(timeout)
            
            if thread.is_alive():
                raise RuntimeError(f"进程 {rank} 初始化分布式环境超时（{timeout}秒），可能存在网络或端口冲突")
            
            # 检查是否初始化成功
            if result_queue.get():
                # 设置GPU设备
                if torch.cuda.is_available():
                    torch.cuda.set_device(rank)
                # 尝试简单的同步检查，使用更可靠的同步方式
                try:
                    # 小批量同步检查，移除timeout参数以兼容旧版PyTorch
                    dist.barrier()
                    print(f"进程 {rank} 同步检查通过")
                except Exception as barrier_error:
                    print(f"进程 {rank} 同步检查失败: {str(barrier_error)}")
                    print(f"进程 {rank} 警告: 同步检查失败但继续执行，后续训练可能不稳定")
                    # 同步检查失败但继续执行，同时记录警告
                
                elapsed = time.time() - start_time
                print(f"进程 {rank} 分布式环境初始化成功，耗时: {elapsed:.2f}秒")
                print(f"进程 {rank} 分布式训练环境就绪: rank={rank}, world_size={world_size}, backend={backend}")
                
                return backend
            else:
                # 获取错误信息并抛出
                error = error_queue.get()
                raise RuntimeError(f"进程 {rank} 初始化失败: {str(error)}") from error
            
        except Exception as e:
            elapsed = time.time() - start_time
            print(f"进程 {rank} 初始化分布式环境失败，耗时: {elapsed:.2f}秒: {str(e)}")
            # 确保即使失败也不会让其他进程卡住

            raise
    
    @staticmethod
    def cleanup():
        """清理分布式环境"""
        try:
            if dist.is_initialized():
                dist.destroy_process_group()
                print("分布式进程组已清理")
        except Exception as e:
            print(f"清理分布式环境时出错: {str(e)}")
    
    @staticmethod
    def wrap_model(model, device_id):
        """将模型包装为DDP模型
        
        Args:
            model: 要包装的PyTorch模型
            device_id: 当前设备ID
            
        Returns:
            DDP: 包装后的DDP模型
        """
        # torch 1.10.2 兼容的DDP包装方式
        return DDP(
            model,
            device_ids=[device_id],
            output_device=device_id,
            # 兼容torch 1.10.2，使用find_unused_parameters=False提高性能
            find_unused_parameters=False
        )
    

    
    @staticmethod
    def get_rank():
        """获取当前进程的rank
        
        Returns:
            int: 当前进程的rank，如果未初始化返回0
        """
        if not dist.is_initialized():
            return 0
        return dist.get_rank()
    



# 添加工具函数用于降级处理
def safe_ddp_setup(rank, world_size):
    """安全的DDP设置函数，失败时返回降级信息
    
    Args:
        rank: 进程排名
        world_size: 世界大小
        
    Returns:
        tuple: (是否成功, 使用的后端, 错误信息)
    """
    try:
        backend = DDPUtil.setup(rank, world_size)
        return True, backend, None
    except Exception as e:
        return False, None, str(e)


def safe_ddp_cleanup():
    """安全的DDP清理函数"""
    try:
        DDPUtil.cleanup()
        return True
    except Exception:
        return False


def is_distributed_env():
    """
    检查是否在分布式训练环境中
    
    Returns:
        bool: 如果在分布式环境中返回True，否则返回False
    """
    return 'LOCAL_RANK' in os.environ






